from app import Strategy, Context, ResultAnalyse

class TzcStrategy(Strategy.Strategy):
    def __init__(self, name, params, context:Context.Context):
        super().__init__(name, params, context)
        self.klines = []
        self.order_id = 1
            
    def on_kline_day(self, data:Context.Kline):
        self.klines.append(data)
        if self.context.has_position(data.code):
            current_position = self.context.account.position[data.code]
            current_position.last_price = data.close
            self.context.account.position[data.code] = current_position
        if self.__will_raise__():
            if self.context.has_position(data.code):
                last_item = self.klines[-1]
                if last_item.close == last_item.high:
                    self.order_id += 1
                    order = Context.Order(
                        order_id=self.order_id,
                        code=data.code,
                        market = data.market,
                        direction=1,
                        price=data.close,
                        qty=1000,
                        date=data.date,
                        time=data.time
                    )
                    self.context.order(order)
            else:
                # print(self.klines[-1].close)
                self.order_id += 1
                order = Context.Order(
                    order_id=self.order_id,
                    code=data.code,
                    market = data.market,
                    direction=1,
                    price=data.close,
                    qty=1000,
                    date=data.date,
                    time=data.time
                )
                self.context.order(order)
        elif self.__will_fall__() and self.context.has_position(data.code):
            current_position =self.context.get_position(data.code)
            self.order_id += 1
            order = Context.Order(
                order_id=self.order_id,
                code=data.code,
                market = data.market,
                direction=0,
                price=data.close,
                qty=current_position.qty,
                date=data.date,
                time=data.time
            )
            self.context.order(order)
        if len(self.klines) > 2:
            self.klines = self.klines[-2:]
        if self.context.account.date != data.date:
            self.context.account.date = data.date
    
    def on_daily_close(self):
        current_asset = 0
        for key,value in self.context.account.position.items():
            current_asset += value.qty * value.last_price
        current_asset += self.context.account.money
        self.context.account.assets[self.context.account.date] = current_asset
        print(f"{self.context.account.date}, {current_asset}")
        print(self.context.account.position)

    def on_complete(self):
        # print(self.context.account)
        analyse = ResultAnalyse.ResultAnalyse()
        analyse.draw_result_html(self.context.account)

    def __will_fall__(self):
        if len(self.klines) != 3:
            return False
        last_item = self.klines[-1]
        last_item2 = self.klines[-2]
        if last_item.low < last_item2.low:
            return True
        return False
            
    def __will_raise__(self):
        if len(self.klines) != 3:
            return False
        last_item = self.klines[-1]
        last_item2 = self.klines[-2]
        last_item3 = self.klines[-3]
        if last_item.close <= last_item2.close or last_item2.close <= last_item3.close:
            return False
        # 最高价递增
        # if last_item.high <= last_item2.high:
        if last_item.high <= last_item2.high or last_item2.high <= last_item3.high:
            return False
        # if last_item.low <= last_item2.low or last_item2.low <= last_item3.low:
        if last_item.low <= last_item2.low:
            return False
        if last_item.close <= last_item2.high:
            return False
        # 今日行情数据进行比较
        if last_item.close <= last_item.open or last_item.close_ema5 < last_item.close_ema10:
            return False
        if 3*(last_item.high - last_item.close) >= last_item.close - last_item.low:
            return False
        # 成交量
        if last_item.volume < last_item2.volume:
            return False
        if last_item.amount / last_item.volume <= last_item.close:
            return False
        return True